package com.alibaba.cloud.ai.example.graph.basic_chatbot;

import lombok.extern.slf4j.Slf4j;
import org.bsc.langgraph4j.GraphStateException;
import org.bsc.langgraph4j.StateGraph;
import org.bsc.langgraph4j.serializer.StateSerializer;
import org.bsc.langgraph4j.state.AgentState;
import org.springframework.stereotype.Service;

import java.util.Map;
import java.util.Optional;

import static org.bsc.langgraph4j.StateGraph.END;
import static org.bsc.langgraph4j.StateGraph.START;
import static org.bsc.langgraph4j.action.AsyncNodeAction.node_async;

/**
 * @author ming（fuyou.lxm）
 * @version 1.0
 * @date 2025/1/20 10:01
 */

@Slf4j
@Service
public class BasicChatbot {

    private final BasicChatBotService basicChatBotService;

    public BasicChatbot(BasicChatBotService basicChatBotService) {
        this.basicChatBotService = basicChatBotService;
    }

    public class GraphBuilder {

        private StateSerializer<BasicChatbot.State> stateSerializer;

        /**
         * Sets the state serializer for the graph builder.
         *
         * @param stateSerializer the state serializer to set
         * @return the current instance of GraphBuilder for method chaining
         */
        public BasicChatbot.GraphBuilder stateSerializer(StateSerializer<BasicChatbot.State> stateSerializer) {
            this.stateSerializer = stateSerializer;
            return this;
        }

        public StateGraph<BasicChatbot.State> build() throws GraphStateException {
            if (stateSerializer == null) {
                stateSerializer = new MemoryStateSerializer<>(BasicChatbot.State::new) {
                };
            }

            return new StateGraph<>(stateSerializer)
                    .addEdge(START, "chatbot")
                    .addNode("chatbot", node_async(this::callAgent))
                    .addEdge("chatbot", END);

        }

        Map<String, Object> callAgent(BasicChatbot.State state) {
            log.info("callAgent");

            var input = state.input()
                    .orElseThrow(() -> new IllegalArgumentException("no input provided!"));

            var response = basicChatBotService.execute(input);

            var finish = new BasicChatbot.Finish(Map.of("returnValues", response));
            return Map.of(BasicChatbot.State.AGENT_OUTCOME, new BasicChatbot.Outcome(finish));
        }
    }

    public final BasicChatbot.GraphBuilder graphBuilder() {
        return new BasicChatbot.GraphBuilder();
    }


    public static class State extends AgentState {
        public static final String INPUT = "input";

        public static final String AGENT_OUTCOME = "outcome";

        /**
         * Constructs a new State object using the initial data provided in the initData map.
         *
         * @param initData the map containing the initial settings for this state
         */
        public State(Map<String, Object> initData) {
            super(initData);
        }

        public Optional<String> input() {
            return value(INPUT);
        }

        public Optional<Outcome> agentOutcome() {
            return value(AGENT_OUTCOME);
        }
    }

    public record Outcome(BasicChatbot.Finish finish) {
    }

    public record Finish(Map<String, Object> returnValues) {
    }
}
